#!/usr/bin/env python3
"""
测试数据增强功能的简单脚本
"""

import cv2
import numpy as np
import os
from augmentation import DataAugmentation

def create_test_data():
    """创建测试用的图像和标签数据"""
    # 创建一个简单的测试图像 (100x100x3)
    img = np.zeros((100, 100, 3), dtype=np.uint8)
    img[20:80, 20:80] = [100, 150, 200]  # 中间的矩形区域
    img[40:60, 40:60] = [255, 255, 255]  # 小的白色方块
    
    # 创建对应的边缘标签 (100x100)
    gt = np.zeros((100, 100), dtype=np.uint8)
    gt[20:21, 20:80] = 255  # 上边缘
    gt[79:80, 20:80] = 255  # 下边缘
    gt[20:80, 20:21] = 255  # 左边缘
    gt[20:80, 79:80] = 255  # 右边缘
    gt[40:41, 40:60] = 255  # 内部方块上边缘
    gt[59:60, 40:60] = 255  # 内部方块下边缘
    gt[40:60, 40:41] = 255  # 内部方块左边缘
    gt[40:60, 59:60] = 255  # 内部方块右边缘
    
    return img, gt

def save_test_results(img, gt, suffix=""):
    """保存测试结果"""
    results_dir = "augmentation_test_results"
    os.makedirs(results_dir, exist_ok=True)
    
    # 保存图像
    cv2.imwrite(os.path.join(results_dir, f"test_img{suffix}.png"), img)
    
    # 保存边缘标签
    cv2.imwrite(os.path.join(results_dir, f"test_gt{suffix}.png"), gt)
    
    print(f"保存结果: test_img{suffix}.png, test_gt{suffix}.png")

def test_individual_augmentations():
    """测试各种增强方法"""
    print("=== 测试各种数据增强方法 ===")
    
    # 创建测试数据
    img, gt = create_test_data()
    save_test_results(img, gt, "_original")
    
    # 测试各种增强
    augmentor = DataAugmentation(
        enable_flip=True,
        enable_rotation=True,
        enable_brightness=True,
        enable_contrast=True,
        enable_noise=True,
        flip_prob=1.0,          # 确保应用
        rotation_prob=1.0,      # 确保应用
        brightness_prob=1.0,    # 确保应用
        contrast_prob=1.0,      # 确保应用
        noise_prob=1.0,         # 确保应用
        rotation_range=30.0,
        brightness_range=0.3,
        contrast_range=0.3,
        noise_std=15.0
    )
    
    print(f"增强器配置: {augmentor}")
    
    # 1. 测试水平翻转
    img_flip_h, gt_flip_h = augmentor.horizontal_flip(img, gt)
    save_test_results(img_flip_h, gt_flip_h, "_flip_horizontal")
    
    # 2. 测试垂直翻转  
    img_flip_v, gt_flip_v = augmentor.vertical_flip(img, gt)
    save_test_results(img_flip_v, gt_flip_v, "_flip_vertical")
    
    # 3. 测试旋转
    img_rot, gt_rot = augmentor.rotate(img, gt, 45.0)
    save_test_results(img_rot, gt_rot, "_rotate_45")
    
    # 4. 测试亮度调整
    img_bright = augmentor.adjust_brightness(img, 0.3)
    save_test_results(img_bright, gt, "_brightness_up")
    
    img_dark = augmentor.adjust_brightness(img, -0.3)
    save_test_results(img_dark, gt, "_brightness_down")
    
    # 5. 测试对比度调整
    img_contrast_up = augmentor.adjust_contrast(img, 1.5)
    save_test_results(img_contrast_up, gt, "_contrast_up")
    
    img_contrast_down = augmentor.adjust_contrast(img, 0.5)
    save_test_results(img_contrast_down, gt, "_contrast_down")
    
    # 6. 测试噪声
    img_noise = augmentor.add_gaussian_noise(img, 20.0)
    save_test_results(img_noise, gt, "_noise")
    
    print("单独增强测试完成！")

def test_random_augmentation():
    """测试随机增强组合"""
    print("\n=== 测试随机增强组合 ===")
    
    # 创建测试数据
    img, gt = create_test_data()
    
    # 创建增强器
    augmentor = DataAugmentation(
        enable_flip=True,
        enable_rotation=True,
        enable_brightness=True,
        enable_contrast=True,
        enable_noise=True,
        flip_prob=0.5,
        rotation_prob=0.5,
        brightness_prob=0.5,
        contrast_prob=0.5,
        noise_prob=0.5
    )
    
    # 生成多个随机增强结果
    for i in range(5):
        img_aug, gt_aug = augmentor.random_augment(img, gt)
        save_test_results(img_aug, gt_aug, f"_random_{i+1}")
    
    print("随机增强测试完成！")

def test_with_dataset_format():
    """测试符合数据集格式的增强"""
    print("\n=== 测试数据集格式增强 ===")
    
    # 模拟从数据集加载的数据格式
    img = np.random.randint(0, 256, (352, 352, 3), dtype=np.uint8)
    gt = np.random.randint(0, 2, (352, 352), dtype=np.uint8) * 255
    
    # 创建增强器（使用更保守的参数）
    augmentor = DataAugmentation(
        enable_flip=True,
        enable_rotation=True,
        enable_brightness=True,
        enable_contrast=True,
        enable_noise=True,
        flip_prob=0.5,
        rotation_prob=0.3,
        brightness_prob=0.3,
        contrast_prob=0.3,
        noise_prob=0.2,
        rotation_range=15.0,
        brightness_range=0.2,
        contrast_range=0.2,
        noise_std=10.0
    )
    
    # 应用增强
    img_aug, gt_aug = augmentor.random_augment(img, gt)
    
    # 保存结果
    save_test_results(img_aug, gt_aug, "_dataset_format")
    
    print(f"原始图像形状: {img.shape}, 增强后图像形状: {img_aug.shape}")
    print(f"原始标签形状: {gt.shape}, 增强后标签形状: {gt_aug.shape}")
    print("数据集格式测试完成！")

if __name__ == "__main__":
    print("开始测试数据增强功能...")
    
    try:
        # 运行各种测试
        test_individual_augmentations()
        test_random_augmentation()
        test_with_dataset_format()
        
        print("\n✓ 所有测试完成！请查看 'augmentation_test_results' 目录中的结果图像。")
        
    except Exception as e:
        print(f"\n✗ 测试失败: {str(e)}")
        import traceback
        traceback.print_exc()